{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "# from sklearn import datasets\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import TensorDataset\n",
    "\n",
    "\n",
    "class KddData(object):\n",
    "\n",
    "    def __init__(self, batch_size, file_name1, file_name2):\n",
    "        data1 = pd.read_csv(file_name1)\n",
    "        data2 = pd.read_csv(file_name2)\n",
    "        \n",
    "        \n",
    "        self._encoder = {\n",
    "            'label':    LabelEncoder()\n",
    "        }\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        target = np.array(data1['label'])\n",
    "        features = np.array(data1.drop('label', axis=1))\n",
    "        data_X, data_y = self.__encode_data(features, target)\n",
    "        self.train_dataset = TensorDataset(\n",
    "            torch.from_numpy(data_X.astype(np.float32)),\n",
    "            torch.from_numpy(data_y.astype(np.int64))\n",
    "        )\n",
    "\n",
    "        target = np.array(data2['label'])\n",
    "        features = np.array(data2.drop('label', axis=1))\n",
    "        data_X, data_y = self.__encode_data(features, target)\n",
    "        self.test_dataset = TensorDataset(\n",
    "            torch.from_numpy(data_X.astype(np.float32)),\n",
    "            torch.from_numpy(data_y.astype(np.int64))\n",
    "        )\n",
    "\n",
    " \n",
    "        self.train_dataloader = DataLoader(self.train_dataset, self.batch_size, shuffle=True)\n",
    "        self.test_dataloader = DataLoader(self.test_dataset, self.batch_size, shuffle=True)\n",
    "\n",
    "    \"\"\"将数据中字符串部分转换为数字，并将输入的41维特征转换为8*8的矩阵\"\"\"\n",
    "    def __encode_data(self, data_X, data_y):\n",
    "        self._encoder['label'].fit(list(set(data_y)))\n",
    "        data_X = np.pad(data_X, ((0, 0), (0, 100 - len(data_X[0]))), 'constant').reshape(-1, 1, 100)\n",
    "        data_y = self._encoder['label'].transform(data_y)\n",
    "        return data_X, data_y\n",
    "\n",
    "    \"\"\"将数据拆分为训练集和测试集，并转换为TensorDataset对象\"\"\"\n",
    "    def __split_data_to_tensor(self, data_X, data_y):\n",
    "        X_train, X_test, y_train, y_test = train_test_split(data_X, data_y, test_size=0.3)\n",
    "        train_dataset = TensorDataset(\n",
    "            torch.from_numpy(X_train.astype(np.float32)),\n",
    "            torch.from_numpy(y_train.astype(np.int64))\n",
    "        )\n",
    "        test_dataset = TensorDataset(\n",
    "            torch.from_numpy(X_test.astype(np.float32)),\n",
    "            torch.from_numpy(y_test.astype(np.int64))\n",
    "        )\n",
    "        return train_dataset, test_dataset\n",
    "\n",
    "    \"\"\"接受一个数组进行解码\"\"\"\n",
    "    def decode(self, data, label=False):\n",
    "        if not label:\n",
    "            _data = list(data)\n",
    "            _data[1] = self._encoder['protocal'].inverse_transform([_data[1]])[0]\n",
    "            _data[2] = self._encoder['service'].inverse_transform([_data[2]])[0]\n",
    "            _data[2] = self._encoder['flag'].inverse_transform([_data[3]])[0]\n",
    "            return _data\n",
    "        return self._encoder['label'].inverse_transform(data)\n",
    "    \n",
    "    def encode(self, data, label=False):\n",
    "        if not label:\n",
    "            _data = list(data)\n",
    "            _data[1] = self._encoder['protocal'].transform([_data[1]])[0]\n",
    "            _data[2] = self._encoder['service'].transform([_data[2]])[0]\n",
    "            _data[3] = self._encoder['flag'].transform([_data[3]])[0]\n",
    "            return _data\n",
    "        return self._encoder['label'].transform([data])[0]\n",
    "\n",
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self, in_dim, num_class):\n",
    "        super(CNN, self).__init__()\n",
    "\n",
    "        self.conv1 = nn.Sequential(       \n",
    "            nn.Conv1d(1, 32, 5, 1, 2),\n",
    "            nn.BatchNorm1d(32), \n",
    "            nn.ELU(),                     \n",
    "            nn.Conv1d(32, 32, 5, 1, 2),\n",
    "            nn.BatchNorm1d(32), \n",
    "            nn.ELU(),           \n",
    "            nn.MaxPool1d(3, 3, 0),\n",
    "            nn.Dropout(0.1)\n",
    "        )\n",
    "        self.conv2 = nn.Sequential(       \n",
    "            nn.Conv1d(32, 64, 5, 1, 2),\n",
    "            nn.BatchNorm1d(64), \n",
    "            nn.ReLU(),                     \n",
    "            nn.Conv1d(64, 64, 5, 1, 2),\n",
    "            nn.BatchNorm1d(64), \n",
    "            nn.ReLU(),           \n",
    "            nn.MaxPool1d(3, 3, 0),\n",
    "            nn.Dropout(0.1)\n",
    "        )\n",
    "        self.out1 = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(704, 512),\n",
    "            nn.BatchNorm1d(512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 256),\n",
    "            nn.BatchNorm1d(256),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.1)\n",
    "        )\n",
    "        self.out2 = nn.Sequential(\n",
    "            nn.Linear(256, num_class)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        output = self.out1(x)\n",
    "        output = self.out2(output)\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "# from ignite.metrics import Precision, Recall,\n",
    "\n",
    "# 神经网络参数\n",
    "batch_size = 32\n",
    "learning_rate = 1e-2\n",
    "num_epoches = 40\n",
    "USE_GPU = False\n",
    "num_class = 2\n",
    "\n",
    "trainFile = 'train-wired.csv'\n",
    "testFile = 'test-wired.csv'\n",
    "testFile2 = 'test-kr.csv'\n",
    "dataset = KddData(batch_size, trainFile, testFile2)\n",
    "model = CNN(1, num_class)\n",
    "\n",
    "def train():\n",
    "    \n",
    "    global model\n",
    "\n",
    "    if USE_GPU:\n",
    "        model = model.cuda()\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "    for epoch in range(num_epoches):\n",
    "        print('epoch {}'.format(epoch + 1))\n",
    "        print('*' * 10)\n",
    "        running_loss = 0.0\n",
    "        running_acc = 0.0\n",
    "        for i, data in enumerate(dataset.train_dataloader, 1):\n",
    "            img, label = data\n",
    "            if USE_GPU:\n",
    "                img = img.cuda()\n",
    "                label = label.cuda()\n",
    "            img = Variable(img)\n",
    "            label = Variable(label)\n",
    "            # 向前传播\n",
    "            out = model(img)\n",
    "            loss = criterion(out, label)\n",
    "            running_loss += loss.item() * label.size(0)\n",
    "            _, pred = torch.max(out, 1)\n",
    "            num_correct = (pred == label).sum()\n",
    "            accuracy = (pred == label).float().mean()\n",
    "            running_acc += num_correct.item()\n",
    "            # 向后传播\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(\n",
    "            epoch + 1, running_loss / (len(dataset.train_dataset)), running_acc / (len(\n",
    "                dataset.train_dataset))))\n",
    "        model.eval()\n",
    "            \n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5 (default, Sep  3 2020, 21:29:08) [MSC v.1916 64 bit (AMD64)]"
  },
  "metadata": {
   "interpreter": {
    "hash": "11f1dc213e07634baa4c5c321dec03c05dafae643c50f20e6d1a492290c05dc2"
   }
  },
  "vscode": {
   "interpreter": {
    "hash": "e42634819b8c191a5d07eaf23810ff32516dd8d3875f28ec3e488928fbd3c187"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
